from python_ai.common.xcommon import *
import tensorflow.compat.v1 as tf
import tensorflow as tsf
import numpy as np
import sys

tf.set_random_seed(777)

n_input = 3
n_hidden = 5
n_steps = 6

X0_batch = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 0, 1]])  #(4,3)
X1_batch = np.array([[9, 8, 7], [0, 0, 0], [6, 5, 4], [3, 2, 1]])
X2_batch = X0_batch.copy()
X3_batch = X1_batch.copy()
X4_batch = X0_batch.copy()
X5_batch = X1_batch.copy()
X_batch = np.array([X0_batch, X1_batch, X2_batch, X3_batch,
             X4_batch, X5_batch])  # 6x4x3
print(X_batch.shape)
X_batch = np.transpose(X_batch, [1, 0, 2])  # 4x6x3
print(X_batch.shape)
# sys.exit(0)

X = tf.placeholder(tf.float32, [None, n_steps, n_input], 'X')
print(f'X: {X}')

cell = tf.nn.rnn_cell.BasicRNNCell(num_units=n_hidden)

output_seqs, states = tf.nn.dynamic_rnn(cell, X, dtype=tf.float32)
print(f'output_seqs: {output_seqs}')
print(f'states: {states}')

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    output_seqs_, states_ = sess.run([output_seqs, states],
                                           feed_dict={
                                               X: X_batch,
                                           })

for i, output in enumerate(output_seqs_):
    print(f'#{i}: {output}')
print('states:')
print(states_)
print(f'outputs: {np.shape(output_seqs_)}')
print(f'states: {np.shape(states_)}')
